Write a custom CUDA kernel to optimize `Large-Margin Softmax Loss` (L-Softmax).

Formula for target class y_i:
f_yi = |W_yi| * |x_i| * psi(theta_yi)
Where psi(theta) = (-1)^k * cos(m * theta) - 2k, for theta in [k*pi/m, (k+1)*pi/m].
Non-target classes f_j remain |W_j| * |x_i| * cos(theta_j) (standard dot product).

Problem Analysis:
1. Complex Piecewise Function: Computing psi(theta) involves `acos`, `floor`, `cos`, and conditional logic dependent on the angle region `k`. Implementing this via vector masks in Python is inefficient and memory-heavy.
2. Operator Chaining: The sequence `div` (get cos) -> `acos` -> `compute psi` -> `mul` (restore scale) -> `softmax` -> `nll` involves multiple passes.

Optimization Strategy: Fused L-Softmax Kernel

1. Inputs:
   - `logits`: Raw output of FC layer (N, C).
   - `x_norm`: Norm of input features (N,).
   - `w_norm`: Norm of weights (C,).
   - `targets`: Ground truth labels (N).
   - `m`: Margin integer.
   - `lambda`: Weight for gradual annealing (optional, but standard in L-Softmax implementation: `(f_yi + lambda*original) / (1+lambda)`). Let's implement the pure version or simple weighted version.

2. One-Block-per-Row: Assign one CUDA block to process one sample.

3. Fused Logic:
   - Load `logits` into Shared Memory (float4).
   - Thread 0 (or responsible thread) identifies `target`.
   - Retrieve `logits[target]`, `x_norm[row]`, `w_norm[target]`.
   - Compute `cos_theta = logits[target] / (x_norm * w_norm)`.
   - Compute `psi_theta` using the piecewise formula with fast math intrinsics.
   - Update `logits[target] = x_norm * w_norm * psi_theta`.
   - Perform standard Online Softmax (Max + SumExp) reductions.
   - Compute NLL Loss. 
  
Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is:   
  
```python
import torch
import torch.nn as nn
import torch.nn.functional as F
import math

BATCH_SIZE = 128
NUM_CLASSES = 1000
FEAT_DIM = 512
SHAPE_X = (BATCH_SIZE, FEAT_DIM)

MARGIN_M = 2
LAMBDA_VAL = 0.0 
REDUCTION = 'none'

class LSoftmaxLoss(nn.Module):
    '''
    Large-Margin Softmax Loss for Convolutional Neural Networks
    https://arxiv.org/pdf/1904.11138
    '''
    def __init__(self, num_classes, feat_dim, margin=2, lambda_val=0.0, reduction='mean', weight=None):
        super(LSoftmaxLoss, self).__init__()
        self.margin = margin
        self.lambda_val = lambda_val
        self.reduction = reduction
        self.pi = math.pi
        
        if weight is None:
            self.weight = nn.Parameter(torch.randn(num_classes, feat_dim))
            nn.init.xavier_normal_(self.weight)
        else:
            # weight 可学习参数
            self.weight = nn.Parameter(weight)

    def forward(self, x: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        # Compute Logits and Norms
        logits = F.linear(x, self.weight)
        
        x_norm = torch.norm(x, p=2, dim=1) 
        w_norm = torch.norm(self.weight, p=2, dim=1) 
        
        # Get Target Cosine
        target_logits = logits.gather(1, targets.view(-1, 1)).squeeze(1)
        target_w_norm = w_norm[targets]
        
        x_norm_clamped = torch.clamp(x_norm, min=1e-8)
        norm_prod = x_norm_clamped * target_w_norm
        cos_theta = target_logits / norm_prod
        
        cos_theta = cos_theta.clamp(-1.0 + 1e-7, 1.0 - 1e-7)
        
        # Compute Psi
        theta = torch.acos(cos_theta)
        k = (theta * self.margin / self.pi).floor()
        
        minus_one_pow_k = torch.pow(-1, k)
        cos_m_theta = torch.cos(theta * self.margin)
        psi_theta = minus_one_pow_k * cos_m_theta - 2 * k
        
        # Scale back
        target_logits_l = norm_prod * psi_theta
        
        # Apply Lambda
        if self.lambda_val > 0:
            final_target_logit = (target_logits_l + self.lambda_val * target_logits) / (1 + self.lambda_val)
        else:
            final_target_logit = target_logits_l
            
        # Replace in Logits
        one_hot = torch.zeros_like(logits)
        one_hot.scatter_(1, targets.view(-1, 1), 1.0)
        
        output_logits = (1.0 - one_hot) * logits + one_hot * final_target_logit.view(-1, 1)
        
        # Cross Entropy
        loss = F.cross_entropy(output_logits, targets, reduction='none')
        
        if self.reduction == 'mean':
            return loss.mean()
        elif self.reduction == 'sum':
            return loss.sum()
        return loss

class Model(nn.Module):
    def __init__(self, num_classes, feat_dim, margin, lambda_val, reduction, weight):
        super(Model, self).__init__()
        self.loss_fn = LSoftmaxLoss(num_classes, feat_dim, margin, lambda_val, reduction, weight)
    
    def forward(self, x, targets):
        return self.loss_fn(x, targets)

def get_inputs():
    x = torch.randn(SHAPE_X, dtype=torch.float32)
    targets = torch.randint(0, NUM_CLASSES, (BATCH_SIZE,), dtype=torch.long)
    return [x.contiguous(), targets.contiguous()]

def get_init_inputs():
    # 生成正态分布随机数
    weight = torch.randn(NUM_CLASSES, FEAT_DIM, dtype=torch.float32)
    nn.init.xavier_normal_(weight)
    
    return [NUM_CLASSES, FEAT_DIM, MARGIN_M, LAMBDA_VAL, REDUCTION, weight]